import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import joblib
from datetime import datetime, timedelta


class SalesPredictor:
    def __init__(self, model_path='sales_model.pkl'):
        self.model_path = model_path
        self.model = None
        self.cutoff_date = None
        self.max_days = 90  # 商品生命周期90天
        self.forecast_days = 14  # 预测未来14天

    def prepare_data(self, raw_df, cutoff_date=None):
        """
        准备训练数据：
        1. 过滤超过90天的数据
        2. 补充缺失日期销量为0
        3. 生成特征和目标变量
        4. 过滤上市不足90天的商品
        """
        # 确保日期格式为YYYY-MM-DD
        raw_df['date_id'] = pd.to_datetime(raw_df['date_id'], format='%Y-%m-%d')
        raw_df['new_date'] = pd.to_datetime(raw_df['new_date'], format='%Y-%m-%d')

        # 应用截止日期过滤
        if cutoff_date:
            self.cutoff_date = pd.to_datetime(cutoff_date, format='%Y-%m-%d')
            raw_df = raw_df[raw_df['date_id'] <= self.cutoff_date]
        else:
            # 如果没有提供截止日期，使用数据中最大日期
            self.cutoff_date = raw_df['date_id'].max()

        # 生成完整日期范围
        full_data = []
        grouped = raw_df.groupby('spu_code')
        mature_spus = []  # 记录上市时间超过90天的SPU

        for spu, group in grouped:
            launch_date = group['new_date'].min()

            # 计算商品在截止日期时的上市天数
            days_since_launch = (self.cutoff_date - launch_date).days + 1

            # 过滤上市不足90天的商品
            if days_since_launch < self.max_days:
                print(f"SPU {spu} skipped: Only {days_since_launch} days since launch (less than {self.max_days} days)")
                continue

            mature_spus.append(spu)

            # 计算商品生命周期内的日期范围
            date_range = pd.date_range(
                start=launch_date,
                periods=self.max_days,
                freq='D'
            )

            full_dates = pd.DataFrame({'date_id': date_range, 'spu_code': spu})

            # 合并实际销售数据
            merged = pd.merge(full_dates, group, on=['date_id', 'spu_code'], how='left')
            merged['daily_sale_qty'] = merged['daily_sale_qty'].fillna(0)
            merged['new_date'] = launch_date

            full_data.append(merged)

        if not mature_spus:
            print("Warning: No mature SPUs found (all have less than 90 days since launch)")
            return np.array([]), np.array([])

        full_df = pd.concat(full_data, ignore_index=True)
        grouped =full_df.groupby('spu_code')
        # 生成训练样本
        X, y = [], []
        for spu, group in grouped:
            sales = group.sort_values('date_id')['daily_sale_qty'].values

            # 确保有90天的数据（不足补0）
            if len(sales) < self.max_days:
                sales = np.pad(sales, (0, self.max_days - len(sales)), "edge")

            # 生成每个可能的预测点
            for start_day in range(1, self.max_days - self.forecast_days + 1):
                # 特征：历史销量 + 时间特征
                features = list(sales[:start_day])

                # 填充0使特征长度一致
                features += [0] * (self.max_days - len(features))

                # 添加时间特征
                features.append(start_day)  # 当前天数
                features.append(sum(sales[:start_day]))  # 累计销量 到当前日期累计销量
                features.append(np.mean(sales[:start_day]))  # 平均销量 到当前日期平均销量

                # 目标：未来14天累计销量
                end_day = min(start_day + self.forecast_days, len(sales))
                future_sales = sum(sales[start_day:end_day])

                X.append(features)
                y.append(future_sales)

        print(f"Prepared training data from {len(mature_spus)} mature SPUs")
        return np.array(X), np.array(y)

    def full_train(self, data_path, cutoff_date=None):
        """全量训练模型"""
        # 读取Excel文件
        if data_path.endswith('.xlsx'):
            raw_df = pd.read_excel(data_path, engine='openpyxl')
        elif data_path.endswith('.xls'):
            raw_df = pd.read_excel(data_path)
        else:
            raise ValueError("Unsupported file format. Only .xlsx and .xls files are supported.")

        X, y = self.prepare_data(raw_df, cutoff_date)

        # 检查是否有足够的数据
        if len(X) == 0:
            print("Warning: No training data available after filtering.")
            return None

        # 训练模型
        self.model = RandomForestRegressor(n_estimators=100, random_state=42)
        self.model.fit(X, y)

        # 保存模型
        joblib.dump(self.model, self.model_path)
        print(f"Model trained on {len(X)} samples and saved to {self.model_path}")
        return self.model

    def incremental_train(self, new_data_path):
        """增量训练模型"""
        # 加载现有模型
        try:
            self.model = joblib.load(self.model_path)
        except:
            print("No existing model found. Creating new model.")
            self.model = RandomForestRegressor(n_estimators=100, random_state=42)

        # 读取Excel文件
        if new_data_path.endswith('.xlsx'):
            new_df = pd.read_excel(new_data_path, engine='openpyxl')
        elif new_data_path.endswith('.xls'):
            new_df = pd.read_excel(new_data_path)
        else:
            raise ValueError("Unsupported file format. Only .xlsx and .xls files are supported.")

        # 准备增量数据
        X_new, y_new = self.prepare_data(new_df, self.cutoff_date)

        # 检查是否有新数据
        if len(X_new) == 0:
            print("Warning: No new training data available after filtering.")
            return None

        # 增量训练
        self.model.n_estimators += 10  # 增加树的数量
        self.model.fit(X_new, y_new)




        # 保存更新后的模型
        joblib.dump(self.model, self.model_path)
        print(f"Model incrementally trained with {len(X_new)} new samples")
        return self.model

    def predict(self, spu_data):
        """预测单个SPU未来14天销量"""
        # 确保模型已加载
        if self.model is None:
            try:
                self.model = joblib.load(self.model_path)
            except:
                raise ValueError("Model not trained or loaded. Please train a model first.")

        # 确保输入是DataFrame
        if not isinstance(spu_data, pd.DataFrame):
            raise ValueError("Input must be a pandas DataFrame")

        # 检查数据是否为空
        if len(spu_data) == 0:
            print("Warning: Input data is empty.")
            return 0

        # 转换日期格式
        spu_data['date_id'] = pd.to_datetime(spu_data['date_id'], format='%Y-%m-%d')
        spu_data['new_date'] = pd.to_datetime(spu_data['new_date'], format='%Y-%m-%d')

        # 获取上市日期和最后观测日期
        launch_date = spu_data['new_date'].min()
        last_observed_date = spu_data['date_id'].max()
        current_day = (last_observed_date - launch_date).days

        # 检查是否超过预测范围
        if current_day >= self.max_days - self.forecast_days:
            spu_code = spu_data['spu_code'].iloc[0] if 'spu_code' in spu_data.columns else 'Unknown'
            print(
                f"SPU {spu_code} exceeds {self.max_days}-day lifecycle for {self.forecast_days}-day forecast (current day: {current_day}), no prediction made")
            return 0
        if current_day < 0:
            spu_code = spu_data['spu_code'].iloc[0] if 'spu_code' in spu_data.columns else 'Unknown'
            print(
                f"Invalid date range for SPU {spu_code} (launch date: {launch_date}, last date: {last_observed_date})")
            return 0

        # 生成完整90天的日期范围
        full_dates = pd.date_range(start=launch_date, periods=self.max_days, freq='D')
        full_days_df = pd.DataFrame({
            'date_id': full_dates,
            'days_since_launch': range(0, self.max_days)
        })

        # 为原始数据计算天数
        spu_data['days_since_launch'] = (spu_data['date_id'] - launch_date).dt.days

        # 合并完整日期
        merged = pd.merge(full_days_df, spu_data, on=['date_id', 'days_since_launch'], how='left')
        # 填充spu_code和new_date
        merged['spu_code'] = merged['spu_code'].fillna(spu_data['spu_code'].iloc[0])
        merged['new_date'] = launch_date
        merged['daily_sale_qty'] = merged['daily_sale_qty'].fillna(0)

        # 按天数排序
        merged = merged.sort_values('days_since_launch')
        sales = merged['daily_sale_qty'].values[:self.max_days]  # 确保长度为90

        # 构建特征向量
        features = list(sales)  # 长度为90
        features.append(current_day)  # 当前天数
        features.append(sum(sales))  # 累计销量
        features.append(np.mean(sales[:current_day + 1]) if current_day >= 0 else 0)  # 平均销量（只计算到当前日期）

        # 预测
        try:
            prediction = self.model.predict([features])[0]
            return max(0, prediction)  # 确保非负预测
        except Exception as e:
            print(f"Prediction failed: {str(e)}")
            return 0


# 示例用法
if __name__ == "__main__":
    predictor = SalesPredictor()

    # 全量训练（使用截止日期）
    predictor.max_days= 45
    predictor.full_train('sales.xlsx')

    # 准备预测数据示例

    sales_array = [1823,258,115,92,87,76,63,53,66,49,58,52,93,69,62,53,43,34,64,66,71,48,58,58,39,42,69,63,53,62]
    for i in range(len(sales_array)):
        sub_sales_array = sales_array[:i+1]
        known_days = len(sub_sales_array)
        new_spu_data = pd.DataFrame({
            'spu_code': ['LO2412090'] * known_days,
            'date_id': pd.date_range(start='2025-01-10', periods=known_days, freq='D'),
            'new_date': ['2025-01-10'] * known_days,
            'daily_sale_qty': sub_sales_array
        })

        # 执行预测
        forecast = predictor.predict(new_spu_data)
        print(f"上新后第{i+1:02}天销量预测: {forecast:.2f}")